// Copyright 2016 The Draco Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
#ifndef DRACO_CORE_ENCODER_BUFFER_H_
#define DRACO_CORE_ENCODER_BUFFER_H_

#include <memory>
#include <vector>

#include "draco/core/bit_utils.h"
#include "draco/core/macros.h"

namespace draco {

    // Class representing a buffer that can be used for either for byte-aligned
    // encoding of arbitrary data structures or for encoding of variable-length
    // bit data.
    class EncoderBuffer {

        public:

            EncoderBuffer ();
            void Clear ();
            void Resize ( int64_t nbytes );

            // Start encoding a bit sequence. A maximum size of the sequence needs to
            // be known upfront.
            // If encode_size is true, the size of encoded bit sequence is stored before
            // the sequence. Decoder can then use this size to skip over the bit sequence
            // if needed.
            // Returns false on error.
            bool StartBitEncoding ( int64_t required_bits, bool encode_size );

            // End the encoding of the bit sequence and return to the default byte-aligned
            // encoding.
            void EndBitEncoding ();

            // Encode up to 32 bits into the buffer. Can be called only in between
            // StartBitEncoding and EndBitEncoding. Otherwise returns false.
            bool EncodeLeastSignificantBits32 ( int nbits, uint32_t value ) {

                if ( !bit_encoder_active() ) {
                    return false;
                }
                bit_encoder_->PutBits( value, nbits );
                return true;
            }

            // Encode an arbitrary data type.
            // Can be used only when we are not encoding a bit-sequence.
            // Returns false when the value couldn't be encoded.
            template <typename T>
                bool Encode ( const T &data ) {

                    if ( bit_encoder_active() ) {

                        return false;
                    }

                    const uint8_t *src_data = reinterpret_cast<const uint8_t *>( &data );
                    buffer_.insert( buffer_.end(), src_data, src_data + sizeof( T ) );
                    return true;
                }

            bool Encode ( const void *data, size_t data_size ) {

                if ( bit_encoder_active() ) {
                    return false;
                }

                const uint8_t *src_data = reinterpret_cast<const uint8_t *>( data );
                buffer_.insert( buffer_.end(), src_data, src_data + data_size );
                return true;
            }

            bool bit_encoder_active () const {

                return bit_encoder_reserved_bytes_ > 0;
            }

            const char *data () const {

                return buffer_.data();
            }

            size_t size() const {

                return buffer_.size();
            }

            std::vector<char> *buffer () {

                return &buffer_;
            }

        private:
            // Internal helper class to encode bits to a bit buffer.
            class BitEncoder {

                public:
                    // |data| is the buffer to write the bits into.
                    explicit BitEncoder ( char *data ) : bit_buffer_( data ), bit_offset_( 0 ) {}

                    // Write |nbits| of |data| into the bit buffer.
                    void PutBits ( uint32_t data, int32_t nbits ) {

                        DRACO_DCHECK_GE( nbits, 0 );
                        DRACO_DCHECK_LE( nbits, 32 );

                        for ( int32_t bit = 0; bit < nbits; ++bit ) {

                            PutBit( ( data >> bit ) & 1 );
                        }
                    }

                    // Return number of bits encoded so far.
                    uint64_t Bits () const {
                        return static_cast<uint64_t>( bit_offset_ );
                    }

                    // TODO(fgalligan): Remove this function once we know we do not need the
                    // old API anymore.
                    // This is a function of an old API, that currently does nothing.
                    void Flush( int /* left_over_bit_value */ ) {}

                    // Return the number of bits required to store the given number
                    static uint32_t BitsRequired( uint32_t x ) {
                        return static_cast<uint32_t>( MostSignificantBit( x ) );
                    }

                private:
                    void PutBit ( uint8_t value ) {

                        const int byte_size = 8;
                        const uint64_t off = static_cast<uint64_t>( bit_offset_ );
                        const uint64_t byte_offset = off / byte_size;
                        const int bit_shift = off % byte_size;

                        // TODO(fgalligan): Check performance if we add a branch and only do one
                        // memory write if bit_shift is 7. Also try using a temporary variable to
                        // hold the bits before writing to the buffer.

                        bit_buffer_[ byte_offset ] &= ~( 1 << bit_shift );
                        bit_buffer_[ byte_offset ] |= value << bit_shift;
                        bit_offset_++;
                    }

                    char *bit_buffer_;
                    size_t bit_offset_;
            };

            friend class BufferBitCodingTest;
            // All data is stored in this vector.
            std::vector<char> buffer_;

            // Bit encoder is used when encoding variable-length bit data.
            // TODO(ostava): Currently encoder needs to be recreated each time
            // StartBitEncoding method is called. This is not necessary if BitEncoder
            // supported reset function which can easily added but let's leave that for
            // later.
            std::unique_ptr<BitEncoder> bit_encoder_;

            // The number of bytes reserved for bit encoder.
            // Values > 0 indicate we are in the bit encoding mode.
            int64_t bit_encoder_reserved_bytes_;

            // Flag used indicating that we need to store the length of the currently
            // processed bit sequence.
            bool encode_bit_sequence_size_;
    };

}  // namespace draco

#endif  // DRACO_CORE_ENCODER_BUFFER_H_
